from langchain_core.prompts import ChatPromptTemplate
from patagent.llm import chat_llm
from patagent.model import ApiRequest
from patagent.prompt import EXTEACT_PROMPT

# Prompt
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", EXTEACT_PROMPT),
        ("human", "{question}"),
    ]
)

# Chain
extract_chain = prompt | chat_llm.with_structured_output(schema=ApiRequest)